"""
Comprehensive vector field implementation for PDEJulia.jl

Based on py-pde vectorial field implementation with Julia optimizations.
"""

using LinearAlgebra
using StaticArrays
using Statistics
using ..Grids

"""
    VectorField{T<:Real, N, G<:AbstractGrid{N}} <: AbstractDataField

A vector field defined on an N-dimensional grid, containing vectors with N components.

The data is stored as an (N+1)-dimensional array where the last dimension
corresponds to vector components. For a 2D grid, the data shape is (nx, ny, 2),
where data[:, :, 1] is the x-component and data[:, :, 2] is the y-component.

# Fields
- `_grid::G`: The underlying spatial grid
- `_data::Array{T,N+1}`: The field data array with shape (grid_dims..., num_components)

# Examples
```julia
# Create from function
grid = UnitGrid([32, 32])
field = VectorField(pos -> [sin(pos[1]), cos(pos[2])], grid)

# Create from scalar fields
fx = ScalarField(pos -> sin(pos[1]), grid)
fy = ScalarField(pos -> cos(pos[2]), grid)
field = VectorField([fx, fy])

# Create zero field
field = VectorField(Float64, grid)
```
"""
struct VectorField{T<:Real, N, G<:AbstractGrid{N}} <: AbstractDataField
    _grid::G
    _data::Array{T,N+1}  # Shape: (grid_dims..., num_components)
    
    function VectorField{T,N,G}(grid::G, data::Array{T,N+1}) where {T<:Real, N, G<:AbstractGrid{N}}
        expected_size = (size(grid)..., N)
        if size(data) != expected_size
            throw(DimensionMismatch(
                "Data size $(size(data)) does not match expected $(expected_size)"
            ))
        end
        new{T,N,G}(grid, data)
    end
end

# Interface implementation
grid(field::VectorField) = field._grid
data(field::VectorField) = field._data

# Constructors

"""
    VectorField(data::Array{T,N+1}, grid::AbstractGrid{N}) where {T<:Real, N}

Create vector field from data array and grid.
"""
function VectorField(data::Array{T,N+1}, grid::G) where {T<:Real, N, G<:AbstractGrid{N}}
    VectorField{T,N,G}(grid, copy(data))
end

"""
    VectorField(func::Function, grid::AbstractGrid{N}, ::Type{T}=Float64) where {T<:Real, N}

Create vector field by evaluating function at grid points.
The function should take a coordinate vector and return a vector.
"""
function VectorField(func::Function, grid::G, ::Type{T}=Float64) where {T<:Real, N, G<:AbstractGrid{N}}
    data = Array{T,N+1}(undef, size(grid)..., N)
    
    for idx in CartesianIndices(size(grid))
        coord = coordinate(grid, Tuple(idx))
        vec_val = func(coord)
        for i in 1:N
            data[idx, i] = T(vec_val[i])
        end
    end
    
    VectorField{T,N,G}(grid, data)
end

"""
    VectorField(fields::Vector{ScalarField{T}}) where T

Create vector field from a vector of scalar fields (components).
All scalar fields must be on the same grid.
"""
function VectorField(fields::Vector{ScalarField{T}}) where T
    if isempty(fields)
        throw(ArgumentError("Cannot create VectorField from empty vector of ScalarFields"))
    end
    
    g = grid(fields[1])
    N = length(fields)
    
    # Check all grids are compatible
    for field in fields[2:end]
        if grid(field) != g
            throw(ArgumentError("All scalar fields must be on the same grid"))
        end
    end
    
    # Check dimension consistency
    if N != ndim(g)
        throw(DimensionError(
            "Number of scalar fields ($N) must match grid dimension ($(ndim(g)))"
        ))
    end
    
    # Combine data
    data = Array{T,ndim(g)+1}(undef, size(g)..., N)
    for (i, field) in enumerate(fields)
        data[ntuple(i -> i == ndim(g)+1 ? i : Colon(), ndim(g)+1)...] = data(field)
    end
    
    VectorField{T,ndim(g),typeof(g)}(g, data)
end

"""
    VectorField(::Type{T}, grid::AbstractGrid{N}) where {T<:Real, N}

Create zero-initialized vector field of given type.
"""
function VectorField(::Type{T}, grid::G) where {T<:Real, N, G<:AbstractGrid{N}}
    data = zeros(T, size(grid)..., N)
    VectorField{T,N,G}(grid, data)
end

"""
    VectorField(grid::AbstractGrid{N}) where N

Create zero-initialized vector field with Float64 type.
"""
VectorField(grid::AbstractGrid) = VectorField(Float64, grid)# Additional VectorField functionality
# This will be appended to vector.jl

# Base interface
function Base.similar(field::VectorField{T,N,G}) where {T,N,G}
    VectorField{T,N,G}(field._grid, similar(field._data))
end

function Base.copy(field::VectorField{T,N,G}) where {T,N,G}
    VectorField{T,N,G}(field._grid, copy(field._data))
end

Base.size(field::VectorField) = size(field._grid)
Base.length(field::VectorField) = length(field._grid) * ndim(field._grid)
Base.eltype(field::VectorField{T}) where T = T
Base.ndims(field::VectorField{T,N}) where {T,N} = N

# Component access

"""
    getindex(field::VectorField, i::Int) -> ScalarField
    field[i]

Get the i-th component as a scalar field.
"""
function Base.getindex(field::VectorField{T,N,G}, i::Int) where {T,N,G}
    if !(1 ≤ i ≤ N)
        throw(BoundsError(field, i))
    end
    
    component_data = data(field)[ntuple(d -> d == N+1 ? i : Colon(), N+1)...]
    ScalarField{T}(grid(field), component_data)
end

"""
    setindex!(field::VectorField, value, i::Int)
    field[i] = value

Set the i-th component from a scalar field or numeric value.
"""
function Base.setindex!(field::VectorField{T,N}, value::ScalarField{T}, i::Int) where {T,N}
    if !(1 ≤ i ≤ N)
        throw(BoundsError(field, i))
    end
    
    compatible_fields(field[1], value) || throw(ArgumentError("Incompatible grids"))
    
    data(field)[ntuple(d -> d == N+1 ? i : Colon(), N+1)...] .= data(value)
end

function Base.setindex!(field::VectorField{T,N}, value::Number, i::Int) where {T,N}
    if !(1 ≤ i ≤ N)
        throw(BoundsError(field, i))
    end
    
    data(field)[ntuple(d -> d == N+1 ? i : Colon(), N+1)...] .= T(value)
end

# Component access by axis name (for supported grids)
"""
    getindex(field::VectorField, axis::Symbol) -> ScalarField
    field[:x], field[:y], field[:z]

Get component by axis name (if supported by the grid).
"""
function Base.getindex(field::VectorField, axis::Symbol)
    g = grid(field)
    if hasmethod(axis_index, (typeof(g), Symbol))
        i = axis_index(g, axis)
        return field[i]
    else
        throw(ArgumentError("Grid $(typeof(g)) does not support axis names"))
    end
end

function Base.setindex!(field::VectorField, value, axis::Symbol)
    g = grid(field)
    if hasmethod(axis_index, (typeof(g), Symbol))
        i = axis_index(g, axis)
        field[i] = value
    else
        throw(ArgumentError("Grid $(typeof(g)) does not support axis names"))
    end
end

# Vector operations

"""
    dot(v1::VectorField, v2::VectorField) -> ScalarField
    v1 ⋅ v2

Compute the dot product of two vector fields, returning a scalar field.
"""
function LinearAlgebra.dot(v1::VectorField{T1,N,G}, v2::VectorField{T2,N,G}) where {T1,T2,N,G}
    compatible_fields(v1[1], v2[1]) || throw(ArgumentError("Incompatible grids"))
    
    T = promote_type(T1, T2)
    result_data = zeros(T, size(grid(v1)))
    
    for i in 1:N
        result_data .+= data(v1)[ntuple(d -> d == N+1 ? i : Colon(), N+1)...] .* 
                       data(v2)[ntuple(d -> d == N+1 ? i : Colon(), N+1)...]
    end
    
    ScalarField{T}(grid(v1), result_data)
end

"""
    cross(v1::VectorField{T,3}, v2::VectorField{T,3}) -> VectorField{T,3}
    v1 × v2

Compute the cross product of two 3D vector fields.
"""
function LinearAlgebra.cross(v1::VectorField{T1,3,G}, v2::VectorField{T2,3,G}) where {T1,T2,G}
    compatible_fields(v1[1], v2[1]) || throw(ArgumentError("Incompatible grids"))
    
    T = promote_type(T1, T2)
    result_data = Array{T,4}(undef, size(grid(v1))..., 3)
    
    # Cross product: (a × b) = (a2*b3 - a3*b2, a3*b1 - a1*b3, a1*b2 - a2*b1)
    a1, a2, a3 = data(v1)[:,:,:,1], data(v1)[:,:,:,2], data(v1)[:,:,:,3]
    b1, b2, b3 = data(v2)[:,:,:,1], data(v2)[:,:,:,2], data(v2)[:,:,:,3]
    
    result_data[:,:,:,1] .= a2 .* b3 .- a3 .* b2
    result_data[:,:,:,2] .= a3 .* b1 .- a1 .* b3
    result_data[:,:,:,3] .= a1 .* b2 .- a2 .* b1
    
    VectorField{T,3,G}(grid(v1), result_data)
end

"""
    magnitude(field::VectorField) -> ScalarField
    norm(field::VectorField) -> ScalarField

Compute the magnitude (Euclidean norm) of the vector field.
"""
function magnitude(field::VectorField{T,N}) where {T,N}
    result_data = zeros(T, size(grid(field)))
    
    for i in 1:N
        component = data(field)[ntuple(d -> d == N+1 ? i : Colon(), N+1)...]
        result_data .+= component .* component
    end
    
    result_data .= sqrt.(result_data)
    ScalarField{T}(grid(field), result_data)
end

# Alias for magnitude
LinearAlgebra.norm(field::VectorField) = magnitude(field)

"""
    magnitude_squared(field::VectorField) -> ScalarField

Compute the squared magnitude of the vector field (more efficient than magnitude).
"""
function magnitude_squared(field::VectorField{T,N}) where {T,N}
    result_data = zeros(T, size(grid(field)))
    
    for i in 1:N
        component = data(field)[ntuple(d -> d == N+1 ? i : Colon(), N+1)...]
        result_data .+= component .* component
    end
    
    ScalarField{T}(grid(field), result_data)
end

"""
    normalize(field::VectorField) -> VectorField

Normalize the vector field to unit magnitude at each point.
"""
function normalize(field::VectorField{T,N,G}) where {T,N,G}
    mag = magnitude(field)
    result = similar(field)
    
    for i in 1:N
        component_slice = ntuple(d -> d == N+1 ? i : Colon(), N+1)
        data(result)[component_slice...] .= data(field)[component_slice...] ./ data(mag)
    end
    
    result
end

"""
    normalize!(field::VectorField) -> VectorField

Normalize the vector field in-place.
"""
function normalize!(field::VectorField{T,N}) where {T,N}
    mag = magnitude(field)
    
    for i in 1:N
        component_slice = ntuple(d -> d == N+1 ? i : Colon(), N+1)
        data(field)[component_slice...] ./= data(mag)
    end
    
    field
end

# Vector calculus operations

"""
    divergence(field::VectorField, bc=nothing) -> ScalarField
    ∇ ⋅ field

Compute the divergence of the vector field.
For a vector field v = (v₁, v₂, ..., vₙ), the divergence is:
∇ ⋅ v = ∂v₁/∂x₁ + ∂v₂/∂x₂ + ... + ∂vₙ/∂xₙ
"""
function divergence(field::VectorField{T,N}, bc=nothing) where {T,N}
    g = grid(field)
    result_data = zeros(T, size(g))
    
    for i in 1:N
        component_data = data(field)[ntuple(d -> d == N+1 ? i : Colon(), N+1)...]
        grad_component = grid_gradient(g, component_data, i, bc)
        result_data .+= grad_component
    end
    
    ScalarField{T}(g, result_data)
end

"""
    curl(field::VectorField{T,2}, bc=nothing) -> ScalarField

Compute the curl of a 2D vector field, returning a scalar field.
For a 2D vector field v = (vₓ, vᵧ), curl is: ∂vᵧ/∂x - ∂vₓ/∂y
"""
function curl(field::VectorField{T,2}, bc=nothing) where T
    g = grid(field)
    
    vx_data = data(field)[:, :, 1]
    vy_data = data(field)[:, :, 2]
    
    # ∂vᵧ/∂x - ∂vₓ/∂y
    dvydx = grid_gradient(g, vy_data, 1, bc)  # ∂vᵧ/∂x
    dvxdy = grid_gradient(g, vx_data, 2, bc)  # ∂vₓ/∂y
    
    curl_data = dvydx .- dvxdy
    ScalarField{T}(g, curl_data)
end

"""
    curl(field::VectorField{T,3}, bc=nothing) -> VectorField{T,3}

Compute the curl of a 3D vector field, returning a vector field.
For a 3D vector field v = (vₓ, vᵧ, vᵤ), curl is:
∇ × v = (∂vᵤ/∂y - ∂vᵧ/∂z, ∂vₓ/∂z - ∂vᵤ/∂x, ∂vᵧ/∂x - ∂vₓ/∂y)
"""
function curl(field::VectorField{T,3}, bc=nothing) where T
    g = grid(field)
    
    vx_data = data(field)[:, :, :, 1]
    vy_data = data(field)[:, :, :, 2]
    vz_data = data(field)[:, :, :, 3]
    
    curl_data = Array{T,4}(undef, size(g)..., 3)
    
    # ∇ × v = (∂vᵤ/∂y - ∂vᵧ/∂z, ∂vₓ/∂z - ∂vᵤ/∂x, ∂vᵧ/∂x - ∂vₓ/∂y)
    curl_data[:,:,:,1] = grid_gradient(g, vz_data, 2, bc) .- grid_gradient(g, vy_data, 3, bc)
    curl_data[:,:,:,2] = grid_gradient(g, vx_data, 3, bc) .- grid_gradient(g, vz_data, 1, bc)
    curl_data[:,:,:,3] = grid_gradient(g, vy_data, 1, bc) .- grid_gradient(g, vx_data, 2, bc)
    
    VectorField{T,3,typeof(g)}(g, curl_data)
end

"""
    gradient(field::VectorField, bc=nothing) -> TensorField

Compute the gradient of the vector field, returning a rank-2 tensor field.
The (i,j) component of the result is ∂vᵢ/∂xⱼ.
"""
function gradient(field::VectorField{T,N}, bc=nothing) where {T,N}
    g = grid(field)
    grad_data = Array{T,N+2}(undef, size(g)..., N, N)
    
    for i in 1:N  # vector component
        component_data = data(field)[ntuple(d -> d == N+1 ? i : Colon(), N+1)...]
        for j in 1:N  # spatial derivative direction
            grad_component = grid_gradient(g, component_data, j, bc)
            grad_data[ntuple(d -> d == N+1 ? i : (d == N+2 ? j : Colon()), N+2)...] = grad_component
        end
    end
    
    Tensor2Field{T,N,typeof(g)}(g, grad_data)
end

"""
    laplacian(field::VectorField, bc=nothing) -> VectorField

Compute the vector Laplacian of the field.
For each component vᵢ, compute ∇²vᵢ.
"""
function laplacian(field::VectorField{T,N,G}, bc=nothing) where {T,N,G}
    g = grid(field)
    result_data = Array{T,N+1}(undef, size(data(field)))
    
    for i in 1:N
        component_data = data(field)[ntuple(d -> d == N+1 ? i : Colon(), N+1)...]
        lap_component = grid_laplace(g, component_data, bc)
        result_data[ntuple(d -> d == N+1 ? i : Colon(), N+1)...] = lap_component
    end
    
    VectorField{T,N,G}(g, result_data)
end

# Mathematical operations

# Element-wise operations
function Base.:+(v1::VectorField{T1,N,G}, v2::VectorField{T2,N,G}) where {T1,T2,N,G}
    compatible_fields(v1[1], v2[1]) || throw(ArgumentError("Incompatible grids"))
    T = promote_type(T1, T2)
    result = VectorField{T,N,G}(grid(v1), data(v1) .+ data(v2))
end

function Base.:-(v1::VectorField{T1,N,G}, v2::VectorField{T2,N,G}) where {T1,T2,N,G}
    compatible_fields(v1[1], v2[1]) || throw(ArgumentError("Incompatible grids"))
    T = promote_type(T1, T2)
    result = VectorField{T,N,G}(grid(v1), data(v1) .- data(v2))
end

function Base.:*(field::VectorField{T,N,G}, scalar::Number) where {T,N,G}
    S = promote_type(T, typeof(scalar))
    VectorField{S,N,G}(grid(field), data(field) .* scalar)
end

Base.:*(scalar::Number, field::VectorField) = field * scalar

function Base.:/(field::VectorField{T,N,G}, scalar::Number) where {T,N,G}
    S = promote_type(T, typeof(scalar))
    VectorField{S,N,G}(grid(field), data(field) ./ scalar)
end

function Base.:-(field::VectorField{T,N,G}) where {T,N,G}
    VectorField{T,N,G}(grid(field), -data(field))
end

# In-place operations
function Base.:+!(v1::VectorField{T,N,G}, v2::VectorField{T,N,G}) where {T,N,G}
    compatible_fields(v1[1], v2[1]) || throw(ArgumentError("Incompatible grids"))
    data(v1) .+= data(v2)
    v1
end

function Base.:-!(v1::VectorField{T,N,G}, v2::VectorField{T,N,G}) where {T,N,G}
    compatible_fields(v1[1], v2[1]) || throw(ArgumentError("Incompatible grids"))
    data(v1) .-= data(v2)
    v1
end

function Base.:*!(field::VectorField, scalar::Number)
    data(field) .*= scalar
    field
end

function Base.:/!(field::VectorField, scalar::Number)
    data(field) ./= scalar
    field
end

# Scalar field multiplication (component-wise)
function Base.:*(scalar_field::ScalarField{T1}, vector_field::VectorField{T2,N,G}) where {T1,T2,N,G}
    compatible_fields(scalar_field, vector_field[1]) || throw(ArgumentError("Incompatible grids"))
    
    T = promote_type(T1, T2)
    result_data = Array{T,N+1}(undef, size(data(vector_field)))
    
    scalar_data = data(scalar_field)
    for i in 1:N
        result_data[ntuple(d -> d == N+1 ? i : Colon(), N+1)...] = 
            scalar_data .* data(vector_field)[ntuple(d -> d == N+1 ? i : Colon(), N+1)...]
    end
    
    VectorField{T,N,G}(grid(vector_field), result_data)
end

Base.:*(vector_field::VectorField, scalar_field::ScalarField) = scalar_field * vector_field

# Analysis functions

"""
    integrate(field::VectorField) -> Vector

Integrate each component of the vector field over the domain.
"""
function integrate(field::VectorField{T,N}) where {T,N}
    g = grid(field)
    results = Vector{T}(undef, N)
    
    if has_volume_weighting(g)
        # For grids with special volume weighting (spherical, cylindrical)
        for i in 1:N
            component_data = data(field)[ntuple(d -> d == N+1 ? i : Colon(), N+1)...]
            results[i] = integrate(g, component_data)
        end
    else
        # Standard Cartesian integration
        vol_element = prod(spacing(g))
        for i in 1:N
            component_data = data(field)[ntuple(d -> d == N+1 ? i : Colon(), N+1)...]
            results[i] = sum(component_data) * vol_element
        end
    end
    
    results
end

"""
    get_statistics(field::VectorField) -> NamedTuple

Get statistical information about each component of the vector field.
"""
function get_statistics(field::VectorField{T,N}) where {T,N}
    component_stats = []
    magnitude_field = magnitude(field)
    
    for i in 1:N
        comp_data = data(field)[ntuple(d -> d == N+1 ? i : Colon(), N+1)...]
        push!(component_stats, (
            min = minimum(comp_data),
            max = maximum(comp_data),
            mean = mean(comp_data),
            std = std(comp_data)
        ))
    end
    
    (
        components = component_stats,
        magnitude = get_statistics(magnitude_field),
        integral = integrate(field)
    )
end

# Utility functions

"""
    to_scalar(field::VectorField, method::Symbol=:magnitude) -> ScalarField

Convert vector field to scalar field using specified method.
Methods: :magnitude, :magnitude_squared, :x, :y, :z (components), :max, :min
"""
function to_scalar(field::VectorField{T,N}, method::Symbol=:magnitude) where {T,N}
    if method == :magnitude
        return magnitude(field)
    elseif method == :magnitude_squared
        return magnitude_squared(field)
    elseif method in (:x, :y, :z)
        axis_map = Dict(:x => 1, :y => 2, :z => 3)
        component_idx = axis_map[method]
        if component_idx > N
            throw(ArgumentError("Field dimension ($N) too small for component $method"))
        end
        return field[component_idx]
    elseif method == :max
        result_data = data(field)[ntuple(d -> d == N+1 ? 1 : Colon(), N+1)...]
        for i in 2:N
            comp_data = data(field)[ntuple(d -> d == N+1 ? i : Colon(), N+1)...]
            result_data = max.(result_data, comp_data)
        end
        return ScalarField{T}(grid(field), result_data)
    elseif method == :min
        result_data = data(field)[ntuple(d -> d == N+1 ? 1 : Colon(), N+1)...]
        for i in 2:N
            comp_data = data(field)[ntuple(d -> d == N+1 ? i : Colon(), N+1)...]
            result_data = min.(result_data, comp_data)
        end
        return ScalarField{T}(grid(field), result_data)
    else
        throw(ArgumentError("Unknown method: $method"))
    end
end

# Display
function Base.show(io::IO, field::VectorField{T,N}) where {T,N}
    print(io, "VectorField{$T,$N}(")
    print(io, "grid=$(grid(field)), ")
    print(io, "size=$(size(field))")
    print(io, ")")
end

function Base.show(io::IO, ::MIME"text/plain", field::VectorField{T,N}) where {T,N}
    println(io, "VectorField{$T,$N}")
    println(io, "  Grid: $(grid(field))")
    println(io, "  Data size: $(size(data(field)))")
    println(io, "  Components: $N")
    
    # Show statistics for each component
    stats = get_statistics(field)
    for (i, comp_stat) in enumerate(stats.components)
        println(io, "  Component $i: [$(comp_stat.min), $(comp_stat.max)], mean=$(comp_stat.mean)")
    end
    println(io, "  Magnitude: [$(stats.magnitude.min), $(stats.magnitude.max)]")
end